package com.async_export_demo.export;

import com.async_export_demo.config.MyExportConfig;
import com.async_export_demo.model.Task;
import com.async_export_demo.service.TaskService;
import com.async_export_demo.util.SpringContextUtil;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import lombok.Getter;
import lombok.extern.java.Log;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.math.BigDecimal;
import java.math.BigInteger;
import java.math.MathContext;
import java.math.RoundingMode;
import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;

@Component
@Log
public class AsyncExportService {

    private static final ThreadPoolTaskExecutor EXECUTOR = new ThreadPoolTaskExecutor();

    /**
     * 进度，保留两位小数
     */
    private static MathContext MATH_CONTEXT = new MathContext(2, RoundingMode.HALF_UP);

    static {
        //设置核心线程数
        EXECUTOR.setCorePoolSize(2);
        //设置最大线程数
        EXECUTOR.setMaxPoolSize(4);
        //设置线程被回收的空闲时长
        EXECUTOR.setKeepAliveSeconds(6);
        //设置队列容量
        EXECUTOR.setQueueCapacity(2);
        //设置线程前缀
        EXECUTOR.setThreadNamePrefix("export-");
        //设置拒绝策略
        EXECUTOR.setRejectedExecutionHandler(new AbortPolicy());
        //初始化线程池
        EXECUTOR.initialize();
    }

    @Autowired
    private TaskService taskService;

    @Autowired
    private Map<ExportEnum, DataProvider> dataProviderMap;

    @Autowired
    private MyExportConfig myExportConfig;

    // 导出配置类
    public static class ExportConfig {
        private int pageSize = 1000;  // 每页数据量
        private int maxRetry = 3;     // 失败重试次数
        private String fileType = "xlsx"; // 文件类型(csv/xlsx)
        private boolean compress = true; // 是否压缩
    }

    public static class ExportContext {
        private final Map<String, Object> params = new HashMap<>();

        public void put(String key, Object value) {
            params.put(key, value);
        }

        public Object get(String key) {
            return params.get(key);
        }
    }

    /**
     * 因为我想让线程池执行拒绝策略时可以拿到 taskId，以便可以更新状态，所以新定义了一个Runnable
     */
    public static class ExportRunnable implements Runnable {

        @Getter
        private final BigInteger taskId;

        private final Runnable runnable;

        public ExportRunnable(BigInteger taskId, Runnable runnable) {
            this.taskId = taskId;
            this.runnable = runnable;
        }

        @Override
        public void run() {
            runnable.run();
        }

    }

    /**
     * 线程池拒绝策略
     */
    public static class AbortPolicy implements RejectedExecutionHandler {
        /**
         * Creates an {@code AbortPolicy}.
         */
        public AbortPolicy() {
        }

        /**
         * Always throws RejectedExecutionException.
         *
         * @param r the runnable task requested to be executed
         * @param e the executor attempting to execute this task
         * @throws RejectedExecutionException always
         */
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            RejectedExecutionException ex = new RejectedExecutionException("Task " + r.toString() +
                    " rejected from " +
                    e.toString());
            if (r instanceof ExportRunnable) {
                ExportRunnable r2 = (ExportRunnable) r;
                Task entity = new Task();
                entity.setId(r2.getTaskId());
                entity.setStatus("FAILED");
                entity.setErrMsg(ex.getMessage());
                TaskService taskService1 = SpringContextUtil.getApplicationContext().getBean(TaskService.class);
                taskService1.updateExportTask(entity);
            }
            throw ex;
        }
    }

    public BigInteger asyncExport(ExportEnum exportEnum) {
        // 超时时长毫秒
        long timeout = myExportConfig.getTimeoutUnit().toMillis(myExportConfig.getTimeout());
        DataProcessor dataProcessor = new DefaultDataProcessor(myExportConfig, exportEnum.getClz());
        return asyncExport(null, dataProviderMap.get(exportEnum), dataProcessor, timeout);
    }

    public BigInteger asyncExport(ExportConfig config, DataProvider<?> dataProvider, DataProcessor dataProcessor, long timeout) {
        BigInteger taskId = taskService.addExportTask();
        EXECUTOR.execute(new ExportRunnable(taskId, () -> {
            try {
                updateTaskStart(taskId);
                executeExport(taskId, dataProvider, dataProcessor, timeout);
                updateTaskEnd(taskId, dataProcessor.getUrl());
            } catch (Exception e) {
                updateTaskFail(taskId, e.getMessage());
            }
        }));

        /*try {
            updateTaskStart(taskId);
            executeExport(taskId, dataProvider, dataProcessor);
            updateTaskEnd(taskId);
        } catch (Exception e) {
            updateTaskFail(taskId, e.getMessage());
        }*/

        return taskId;
    }

    private <T> void executeExport(BigInteger taskId, DataProvider<T> dataProvider, DataProcessor dataProcessor, long timeout) {
        long start = System.currentTimeMillis();
        long currentPage = 1;
        Page<T> page = null;
        // 新建上下文对象,如需传递参数可以放在里面
        ExportContext context = new ExportContext();
        do {
            // todo 有没其他方法判断是否超时？ScheduledExecutorService?
            // 超时
            if (System.currentTimeMillis() - start > timeout) {
                throw new RejectedExecutionException("Task " + taskId + " rejected,系统执行任务超时");
            }
            page = dataProvider.getPageData(currentPage++, 1000, context);
            if (!Objects.isNull(page) && Objects.nonNull(page.getRecords()) && !page.getRecords().isEmpty()) {
                dataProcessor.processData(page.getRecords());
            }
            // 更新导出进度
            updateTaskProgress(taskId, page);
        } while (page.hasNext());
        dataProcessor.end();
        log.info("executeExport cost:" + (System.currentTimeMillis() - start));
    }

    private void updateTaskFail(BigInteger taskId, String message) {
        Task entity = new Task();
        entity.setId(taskId);
        entity.setStatus("FAILED");
        entity.setErrMsg(message);
        taskService.updateExportTask(entity);
    }

    private void updateTaskStart(BigInteger taskId) {
        Task entity = new Task();
        entity.setId(taskId);
        entity.setStatus("导出中");
        entity.setStartTime(LocalDateTime.now());
        taskService.updateExportTask(entity);
    }

    private void updateTaskEnd(BigInteger taskId, String url) {
        Task entity = new Task();
        entity.setId(taskId);
        entity.setStatus("SUCCESS");
        entity.setErrMsg(null);
        entity.setEndTime(LocalDateTime.now());
        entity.setUrl(url);
        taskService.updateExportTask(entity);
    }

    // todo
    // 暂时将进度更新到数据库，可优化
    private void updateTaskProgress(BigInteger taskId, Page<?> page) {
        Task entity = new Task();
        entity.setId(taskId);
        if (page.getTotal() <= 0) {
            entity.setProgress(new BigDecimal(1));
        } else {
            entity.setProgress(new BigDecimal(page.getCurrent()).divide(new BigDecimal(page.getPages()), MATH_CONTEXT));
        }
        log.info("current:" + page.getCurrent() + " totalPage:" + page.getPages() + " progress " + entity.getProgress().toString());
        taskService.updateExportTask(entity);
    }

}
